# Copyright 2024 Allen Institute for AI

# Copyright 2025 Align-Anything Team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


# import prior

OBJNAV_TYPES_THOR = [
    'AlarmClock',
    'Apple',
    # "BaseballBat",
    'BasketBall',
    'Bed',
    'Bowl',
    'Chair',
    'GarbageCan',
    'HousePlant',
    'Laptop',
    'Mug',
    'Sofa',
    'SprayBottle',
    'Television',
    'Toilet',
    'Vase',
]

PICKUP_TYPES_THOR = [
    'AlarmClock',
    'Apple',
    # "BaseballBat",
    'BasketBall',
    'Bowl',
    'Laptop',
    'Mug',
    'SprayBottle',
    'Vase',
]

RECEPTACLE_TYPES_THOR = [
    'ArmChair',
    'Bed',
    # "Bowl",
    # "Cabinet",
    'Chair',
    'CoffeeTable',
    'CounterTop',
    'Desk',
    'DiningTable',
    # "Drawer",
    'Dresser',
    # "Fridge",
    # "GarbageCan",
    # "Microwave",
    # "Pan",
    # "Plate",
    # "Pot",
    'Shelf',
    # "ShelvingUnit",
    'SideTable',
    'Sofa',
    'Stool',
    'TVStand',
]

ALL_PICKUP_TYPES_THOR = [
    'AlarmClock',
    'AluminumFoil',
    'Apple',
    'AppleSliced',
    # "BaseballBat",
    'Book',
    'Boots',
    'Bottle',
    'Bowl',
    'Box',
    'Bread',
    'BreadSliced',
    'ButterKnife',
    'Candle',
    'CD',
    'CellPhone',
    'Cloth',
    'CreditCard',
    'Cup',
    'DishSponge',
    'Dumbbell',
    'Egg',
    'EggCracked',
    'Fork',
    'HandTowel',
    'Kettle',
    'KeyChain',
    'Knife',
    'Ladle',
    'Laptop',
    'Lettuce',
    'LettuceSliced',
    'Mug',
    'Newspaper',
    'Pan',
    'PaperTowelRoll',
    'Pen',
    'Pencil',
    'PepperShaker',
    'Pillow',
    'Plate',
    'Plunger',
    'Pot',
    'Potato',
    'PotatoSliced',
    'RemoteControl',
    'SaltShaker',
    # "Sandwich", #rose: i think this should be pickupable
    'ScrubBrush',
    'SoapBar',
    'SoapBottle',
    'Spatula',
    'Spoon',
    'SprayBottle',
    'Statue',
    'TableTopDecor',
    'TeddyBear',
    'TennisRacket',
    'TissueBox',
    'ToiletPaper',
    'Tomato',
    'TomatoSliced',
    'Towel',
    'Vase',
    'Watch',
    'WateringCan',
    'WineBottle',
]

AI2THOR_OBJECT_TYPE_TO_WORDNET_SYNSET = {
    'AlarmClock': 'alarm_clock.n.01',
    'AluminumFoil': 'aluminum_foil.n.01',
    'Apple': 'apple.n.01',
    'AppleSliced': 'apple.n.01',
    'ArmChair': 'armchair.n.01',
    'BaseballBat': 'baseball_bat.n.01',
    'BasketBall': 'basketball.n.02',
    'Bathtub': 'bathtub.n.01',
    'BathtubBasin': 'bathtub.n.01',
    'Bed': 'bed.n.01',
    'Blinds': 'blind.n.03',
    'Book': 'book.n.02',
    'Boots': 'boot.n.01',
    'Bottle': 'bottle.n.01',
    'Bowl': 'bowl.n.03',
    'Box': 'carton.n.02',
    'Bread': 'bread.n.01',
    'BreadSliced': 'bread.n.01',
    'ButterKnife': 'butter_knife.n.01',
    'CD': 'compact_disk.n.01',
    'Cabinet': 'cabinet.n.01',
    'Candle': 'candle.n.01',
    'Cart': 'handcart.n.01',
    'CellPhone': 'cellular_telephone.n.01',
    'Chair': 'straight_chair.n.01',
    'Cloth': 'fabric.n.01',
    'ClothesDryer': 'clothes_dryer.n.01',
    'CoffeeMachine': 'coffee_maker.n.01',
    'CoffeeTable': 'coffee_table.n.01',
    'CounterTop': 'countertop.n.01',
    'CreditCard': 'credit_card.n.01',
    'Cup': 'cup.n.01',
    'Curtains': 'curtain.n.01',
    'Desk': 'desk.n.01',
    'DeskLamp': 'table_lamp.n.01',
    'Desktop': 'desktop_computer.n.01',
    'DiningTable': 'dining_table.n.01',
    'DishSponge': 'sponge.n.01',
    'DogBed': 'pad.n.04',
    'Doorframe': 'doorframe.n.01',
    'Doorway': 'doorway.n.01',
    'Drawer': 'drawer.n.01',
    'Dresser': 'chest_of_drawers.n.01',
    'Dumbbell': 'dumbbell.n.01',
    'Egg': 'egg.n.02',
    'EggCracked': 'egg.n.02',
    'Faucet': 'faucet.n.01',
    'Floor': 'floor.n.01',
    'FloorLamp': 'floor_lamp.n.01',
    'Footstool': 'footstool.n.01',
    'Fork': 'fork.n.01',
    'Fridge': 'refrigerator.n.01',
    'GarbageBag': 'bin_liner.n.01',
    'GarbageCan': 'ashcan.n.01',
    'HandTowel': 'hand_towel.n.01',
    'HandTowelHolder': 'towel_rack.n.01',
    'HousePlant': 'houseplant.n.01',
    'Kettle': 'kettle.n.01',
    'KeyChain': 'key_ring.n.01',
    'Knife': 'knife.n.01',
    'Ladle': 'ladle.n.01',
    'Laptop': 'laptop.n.01',
    'LaundryHamper': 'clothes_hamper.n.01',
    'Lettuce': 'lettuce.n.03',
    'LettuceSliced': 'lettuce.n.03',
    'LightSwitch': 'switch.n.01',
    'Microwave': 'microwave.n.02',
    'Mirror': 'mirror.n.01',
    'Mug': 'mug.n.04',
    'Newspaper': 'newspaper.n.03',
    'Ottoman': 'footstool.n.01',
    'Painting': 'painting.n.01',
    'Pan': 'pan.n.01',
    'PaperTowelRoll': 'paper_towel.n.01',
    'Pen': 'pen.n.01',
    'Pencil': 'pencil.n.01',
    'PepperShaker': 'pepper_shaker.n.01',
    'Pillow': 'pillow.n.01',
    'Plate': 'plate.n.04',
    'Plunger': 'plunger.n.03',
    'Poster': 'poster.n.01',
    'Pot': 'pot.n.01',
    'Potato': 'potato.n.01',
    'PotatoSliced': 'potato.n.01',
    'RemoteControl': 'remote_control.n.01',
    'RoomDecor': 'decoration.n.01',
    'Safe': 'safe.n.01',
    'SaltShaker': 'saltshaker.n.01',
    'ScrubBrush': 'scrub_brush.n.01',
    'Shelf': 'shelf.n.01',
    'ShelvingUnit': 'shelf.n.01',
    'ShowerCurtain': 'shower_curtain.n.01',
    'ShowerDoor': 'door.n.01',
    'ShowerGlass': 'door.n.01',
    'ShowerHead': 'showerhead.n.01',
    'SideTable': 'stand.n.04',
    'Sink': 'sink.n.01',
    'SinkBasin': 'sink.n.01',
    'SoapBar': 'bar_soap.n.01',
    'SoapBottle': 'soap_dispenser.n.01',
    'Sofa': 'sofa.n.01',
    'Spatula': 'spatula.n.01',
    'Spoon': 'spoon.n.01',
    'SprayBottle': 'atomizer.n.01',
    'Statue': 'statue.n.01',
    'Stool': 'stool.n.01',
    'StoveBurner': 'burner.n.02',
    'StoveKnob': 'knob.n.02',
    'TVStand': 'stand.n.04',
    'TableTopDecor': 'knickknack.n.01',
    'TeddyBear': 'teddy.n.01',
    'Television': 'television_receiver.n.01',
    'TennisRacket': 'tennis_racket.n.01',
    'TissueBox': 'tissue.n.02',
    'Toaster': 'toaster.n.02',
    'Toilet': 'toilet.n.02',
    'ToiletPaper': 'toilet_tissue.n.01',
    'ToiletPaperHanger': 'hanger.n.02',
    'Tomato': 'tomato.n.01',
    'TomatoSliced': 'tomato.n.01',
    'Towel': 'towel.n.01',
    'TowelHolder': 'towel_rack.n.01',
    'VacuumCleaner': 'vacuum.n.04',
    'Vase': 'vase.n.01',
    'Wall': 'wall.n.01',
    'WashingMachine': 'washer.n.03',
    'Watch': 'watch.n.01',
    'WateringCan': 'watering_can.n.01',
    'Window': 'window.n.01',
    'WineBottle': 'wine_bottle.n.01',
}

AI2THOR_OBJECT_TYPE_TO_MOST_SPECIFIC_WORDNET_LEMMA = {
    'AlarmClock': 'alarm_clock',
    'AluminumFoil': 'aluminum_foil',
    'Apple': 'apple',
    'AppleSliced': 'apple',
    'ArmChair': 'armchair',
    'BaseballBat': 'baseball_bat',
    'BasketBall': 'basketball',
    'Bathtub': 'bathtub',
    'BathtubBasin': 'bathtub',
    'Bed': 'bed',
    'Blinds': 'blind',
    'Book': 'book',
    'Boots': 'boot',
    'Bottle': 'bottle',
    'Bowl': 'bowl',
    'Box': 'box',
    'Bread': 'bread',
    'BreadSliced': 'bread',
    'ButterKnife': 'butter_knife',
    'CD': 'compact_disk',
    'Cabinet': 'cabinet',
    'Candle': 'candle',
    'Cart': 'handcart',
    'CellPhone': 'cellular_telephone',
    'Chair': 'chair',
    'Cloth': 'fabric',
    'ClothesDryer': 'clothes_dryer',
    'CoffeeMachine': 'coffee_maker',
    'CoffeeTable': 'coffee_table',
    'CounterTop': 'countertop',
    'CreditCard': 'credit_card',
    'Cup': 'cup',
    'Curtains': 'curtain',
    'Desk': 'desk',
    'DeskLamp': 'table_lamp',
    'Desktop': 'desktop_computer',
    'DiningTable': 'dining_table',
    'DishSponge': 'sponge',
    'DogBed': 'pad',
    'Doorframe': 'doorframe',
    'Doorway': 'doorway',
    'Drawer': 'drawer',
    'Dresser': 'chest_of_drawers',
    'Dumbbell': 'dumbbell',
    'Egg': 'egg',
    'EggCracked': 'egg',
    'Faucet': 'faucet',
    'Floor': 'flooring',
    'FloorLamp': 'floor_lamp',
    'Footstool': 'footstool',
    'Fork': 'fork',
    'Fridge': 'refrigerator',
    'GarbageBag': 'bin_liner',
    'GarbageCan': 'ashcan',
    'HandTowel': 'hand_towel',
    'HandTowelHolder': 'towel_rack',
    'HousePlant': 'houseplant',
    'Kettle': 'boiler',
    'KeyChain': 'key_ring',
    'Knife': 'knife',
    'Ladle': 'ladle',
    'Laptop': 'laptop',
    'LaundryHamper': 'clothes_hamper',
    'Lettuce': 'lettuce',
    'LettuceSliced': 'lettuce',
    'LightSwitch': 'electric_switch',
    'Microwave': 'microwave_oven',
    'Mirror': 'mirror',
    'Mug': 'mug',
    'Newspaper': 'newspaper',
    'Ottoman': 'pouffe',
    'Painting': 'painting',
    'Pan': 'cooking_pan',
    'PaperTowelRoll': 'paper_towel',
    'Pen': 'pen',
    'Pencil': 'pencil',
    'PepperShaker': 'pepper_shaker',
    'Pillow': 'pillow',
    'Plate': 'plate',
    'Plunger': "plumber's_helper",
    'Poster': 'placard',
    'Pot': 'pot',
    'Potato': 'Irish_potato',
    'PotatoSliced': 'Irish_potato',
    'RemoteControl': 'remote_control',
    'RoomDecor': 'decoration',
    'Safe': 'safe',
    'SaltShaker': 'saltshaker',
    'ScrubBrush': 'scrub_brush',
    'Shelf': 'shelf',
    'ShelvingUnit': 'shelf',
    'ShowerCurtain': 'shower_curtain',
    'ShowerDoor': 'door',
    'ShowerGlass': 'door',
    'ShowerHead': 'showerhead',
    'SideTable': 'stand',
    'Sink': 'sink',
    'SinkBasin': 'sink',
    'SoapBar': 'bar_soap',
    'SoapBottle': 'soap_dispenser',
    'Sofa': 'sofa',
    'Spatula': 'spatula',
    'Spoon': 'spoon',
    'SprayBottle': 'atomizer',
    'Statue': 'statue',
    'Stool': 'stool',
    'StoveBurner': 'burner',
    'StoveKnob': 'knob',
    'TVStand': 'stand',
    'TableTopDecor': 'knickknack',
    'TeddyBear': 'teddy_bear',
    'Television': 'television',
    'TennisRacket': 'tennis_racket',
    'TissueBox': 'tissue_paper',
    'Toaster': 'toaster',
    'Toilet': 'crapper',
    'ToiletPaper': 'toilet_tissue',
    'ToiletPaperHanger': 'hanger',
    'Tomato': 'tomato',
    'TomatoSliced': 'tomato',
    'Towel': 'towel',
    'TowelHolder': 'towel_rack',
    'VacuumCleaner': 'vacuum_cleaner',
    'Vase': 'vase',
    'Wall': 'wall',
    'WashingMachine': 'automatic_washer',
    'Watch': 'watch',
    'WateringCan': 'watering_can',
    'Window': 'window',
    'WineBottle': 'wine_bottle',
}

_cached_bad_asset_ids = None


# def bad_asset_ids():
#     global _cached_bad_asset_ids
#     if _cached_bad_asset_ids is None:
#         _cached_bad_asset_ids = prior.load_dataset(
#             dataset="spoc-data", entity="spoc-robot", revision="spoc-additional-annotation"
#         )
#     return _cached_bad_asset_ids
